{
 "cells": [
  {
   "cell_type": "raw",
   "id": "df7d42b9-58a6-434c-a2d7-0b61142f6d3e",
   "metadata": {},
   "source": [
    "---\n",
    "sidebar_position: 5\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2195672-0cab-4967-ba8a-c6544635547d",
   "metadata": {},
   "source": [
    "# Handle Multiple Retrievers\n",
    "\n",
    "Sometimes, a query analysis technique may allow for selection of which retriever to use. To use this, you will need to add some logic to select the retriever to do. We will show a simple example (using mock data) of how to do that."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4079b57-4369-49c9-b2ad-c809b5408d7e",
   "metadata": {},
   "source": [
    "## Setup\n",
    "#### Install dependencies\n",
    "\n",
    "```{=mdx}\n",
    "import IntegrationInstallTooltip from \"@mdx_components/integration_install_tooltip.mdx\";\n",
    "import Npm2Yarn from \"@theme/Npm2Yarn\";\n",
    "\n",
    "<IntegrationInstallTooltip></IntegrationInstallTooltip>\n",
    "\n",
    "<Npm2Yarn>\n",
    "  @langchain/core @langchain/community @langchain/openai zod chromadb\n",
    "</Npm2Yarn>\n",
    "```\n",
    "\n",
    "#### Set environment variables\n",
    "\n",
    "```\n",
    "OPENAI_API_KEY=your-api-key\n",
    "\n",
    "# Optional, use LangSmith for best-in-class observability\n",
    "LANGSMITH_API_KEY=your-api-key\n",
    "LANGCHAIN_TRACING_V2=true\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c20b48b8-16d7-4089-bc17-f2d240b3935a",
   "metadata": {},
   "source": [
    "### Create Index\n",
    "\n",
    "We will create a vectorstore over fake information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1f621694",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Module: null prototype] {\n",
       "  AdminClient: \u001b[36m[class AdminClient]\u001b[39m,\n",
       "  ChromaClient: \u001b[36m[class ChromaClient]\u001b[39m,\n",
       "  CloudClient: \u001b[36m[class CloudClient extends ChromaClient]\u001b[39m,\n",
       "  CohereEmbeddingFunction: \u001b[36m[class CohereEmbeddingFunction]\u001b[39m,\n",
       "  Collection: \u001b[36m[class Collection]\u001b[39m,\n",
       "  DefaultEmbeddingFunction: \u001b[36m[class _DefaultEmbeddingFunction]\u001b[39m,\n",
       "  GoogleGenerativeAiEmbeddingFunction: \u001b[36m[class _GoogleGenerativeAiEmbeddingFunction]\u001b[39m,\n",
       "  HuggingFaceEmbeddingServerFunction: \u001b[36m[class HuggingFaceEmbeddingServerFunction]\u001b[39m,\n",
       "  IncludeEnum: {\n",
       "    Documents: \u001b[32m\"documents\"\u001b[39m,\n",
       "    Embeddings: \u001b[32m\"embeddings\"\u001b[39m,\n",
       "    Metadatas: \u001b[32m\"metadatas\"\u001b[39m,\n",
       "    Distances: \u001b[32m\"distances\"\u001b[39m\n",
       "  },\n",
       "  JinaEmbeddingFunction: \u001b[36m[class JinaEmbeddingFunction]\u001b[39m,\n",
       "  OpenAIEmbeddingFunction: \u001b[36m[class _OpenAIEmbeddingFunction]\u001b[39m,\n",
       "  TransformersEmbeddingFunction: \u001b[36m[class _TransformersEmbeddingFunction]\u001b[39m\n",
       "}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import { Chroma } from \"@langchain/community/vectorstores/chroma\"\n",
    "import { OpenAIEmbeddings } from \"@langchain/openai\"\n",
    "import \"chromadb\";\n",
    "\n",
    "const texts = [\"Harrison worked at Kensho\"]\n",
    "const embeddings = new OpenAIEmbeddings({ modelName: \"text-embedding-3-small\" })\n",
    "const vectorstore = await Chroma.fromTexts(texts, {}, embeddings, {\n",
    "  collectionName: \"harrison\"\n",
    "})\n",
    "const retrieverHarrison = vectorstore.asRetriever(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cb313c9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "const texts = [\"Ankush worked at Facebook\"]\n",
    "const embeddings = new OpenAIEmbeddings({ modelName: \"text-embedding-3-small\" })\n",
    "const vectorstore = await Chroma.fromTexts(texts, {}, embeddings, {\n",
    "  collectionName: \"ankush\"\n",
    "})\n",
    "const retrieverAnkush = vectorstore.asRetriever(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57396e23-c192-4d97-846b-5eacea4d6b8d",
   "metadata": {},
   "source": [
    "## Query analysis\n",
    "\n",
    "We will use function calling to structure the output. We will let it return multiple queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0b51dd76-820d-41a4-98c8-893f6fe0d1ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import { z } from \"zod\";\n",
    "\n",
    "const searchSchema = z.object({\n",
    "    query: z.string().describe(\"Query to look up\"),\n",
    "    person: z.string().describe(\"Person to look things up for. Should be `HARRISON` or `ANKUSH`.\")\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3c79210",
   "metadata": {},
   "source": [
    "```{=mdx}\n",
    "import ChatModelTabs from \"@theme/ChatModelTabs\";\n",
    "\n",
    "<ChatModelTabs customVarName=\"llm\" />\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "783c03c3-8c72-4f88-9cf4-5829ce6745d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import { ChatPromptTemplate } from \"@langchain/core/prompts\";\n",
    "import { RunnableSequence, RunnablePassthrough } from \"@langchain/core/runnables\";\n",
    "\n",
    "const system = `You have the ability to issue search queries to get information to help answer user information.`\n",
    "const prompt = ChatPromptTemplate.fromMessages(\n",
    "[\n",
    "    [\"system\", system],\n",
    "    [\"human\", \"{question}\"],\n",
    "]\n",
    ")\n",
    "const llmWithTools = llm.withStructuredOutput(searchSchema, {\n",
    "name: \"Search\"\n",
    "})\n",
    "const queryAnalyzer = RunnableSequence.from([\n",
    "    {\n",
    "        question: new RunnablePassthrough(),\n",
    "    },\n",
    "    prompt,\n",
    "    llmWithTools\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9564078",
   "metadata": {},
   "source": [
    "We can see that this allows for routing between retrievers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bc1d3863",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{ query: \u001b[32m\"workplace of Harrison\"\u001b[39m, person: \u001b[32m\"HARRISON\"\u001b[39m }"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await queryAnalyzer.invoke(\"where did Harrison Work\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "af62af17-4f90-4dbd-a8b4-dfff51f1db95",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{ query: \u001b[32m\"Workplace of Ankush\"\u001b[39m, person: \u001b[32m\"ANKUSH\"\u001b[39m }"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await queryAnalyzer.invoke(\"where did ankush Work\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7c65b2f-7881-45fc-a47b-a4eaaf48245f",
   "metadata": {},
   "source": [
    "## Retrieval with query analysis\n",
    "\n",
    "So how would we include this in a chain? We just need some simple logic to select the retriever and pass in the search query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4ec0c7fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "const retrievers = {\n",
    "    HARRISON: retrieverHarrison,\n",
    "    ANKUSH: retrieverAnkush,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8dac7866",
   "metadata": {},
   "outputs": [],
   "source": [
    "import { RunnableConfig, RunnableLambda } from \"@langchain/core/runnables\";\n",
    "\n",
    "const chain = async (question: string, config?: RunnableConfig) => {\n",
    "    const response = await queryAnalyzer.invoke(question, config);\n",
    "    const retriever = retrievers[response.person];\n",
    "    return retriever.invoke(response.query, config);\n",
    "}\n",
    "\n",
    "const customChain = new RunnableLambda({ func: chain });"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "232ad8a7-7990-4066-9228-d35a555f7293",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ Document { pageContent: \u001b[32m\"Harrison worked at Kensho\"\u001b[39m, metadata: {} } ]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await customChain.invoke(\"where did Harrison Work\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "28e14ba5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ Document { pageContent: \u001b[32m\"Ankush worked at Facebook\"\u001b[39m, metadata: {} } ]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await customChain.invoke(\"where did ankush Work\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Deno",
   "language": "typescript",
   "name": "deno"
  },
  "language_info": {
   "file_extension": ".ts",
   "mimetype": "text/x.typescript",
   "name": "typescript",
   "nb_converter": "script",
   "pygments_lexer": "typescript",
   "version": "5.3.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
